import gzip
from io import StringIO
from pathlib import Path

import pandas as pd
from pandas.api.types import is_numeric_dtype

from .PDB_constants import *
from .PDB_dihedral_angle import pdb_dihedral_angle
from .PDB_dotarray_volume import pdb_dotarray_volume
from .PDB_dotarray_surface import pdb_dotarray_surface
from .PDB_csv_sort import pdb_csv_sort
from .utils import timer_s, die
from .typing import *


def load_atom_models(name, fp:Path) -> List[Tuple[int, AtomModel]]:
    # 判断是不是压缩文件
    is_zip = False
    with open(fp, 'rb') as f:
        header = f.read(10)
        is_zip = header.startswith(b'\x1f\x8b\x08') # gzip文件头
    
    if is_zip: open_file = lambda fp: gzip.open(fp, 'rt')
    else:      open_file = lambda fp:      open(fp, 'r')
    try:
        with open_file(fp) as fh:
            line = fh.readline()
            # 判断是不是蛋白质
            if 'DNA' in line or 'RNA' in line or 'NUCLEIC ACID' in line:
                die(f'The PDB {name} is not a PROTEIN but a NUCLEIN')
            else:
                fh.seek(0) # 再从开头读起，以免有些建模的PDB不含开头信息
                lines: List[str] = fh.readlines()
    except:
        die(f'The PDB {name} is UNSUPPORTED file format')

    MODEL  = [i for i in range(len(lines)) if lines[i].startswith('MODEL')]
    ENDMDL = [i for i in range(len(lines)) if lines[i].startswith('ENDMDL')]
    
    if not MODEL and not ENDMDL:                # 只有一个构象
        MDL = [line for line in lines if line.startswith('ATOM')]
    elif MODEL and len(MODEL) == len(ENDMDL):   # 有多个构象
        MDL = {i+1: lines[MODEL[i] : ENDMDL[i]] for i in range(len(MODEL))}
    else:
        die(f'The PDB {name} may be corrupted or contain incomplete MODEL')
    
    a_models = []

    if type(MDL) == list:
        idx = -1 # 标志着只有一个构象
        # 去掉氢原子、末端氧原子、同位置多原子
        a_model = [
            line for line in MDL if
                'OXT' not in line
                    and (line[16] == ' ' or line[16] == 'A')
                    and line[-4] != 'H'
        ]
        a_models.append((idx, a_model))

    elif type(MDL) == dict:
        for idx in MDL: # 这里idx不为-1，标志着超过1个构象
            # 只有ATOM开头的才是标准残基原子
            a_model = [
                line for line in MDL[idx] if
                    line.startswith('ATOM')
                       and 'OXT' not in line
                       and (line[16] == ' ' or line[16] == 'A')
                       and line[-4] != 'H'
            ]
            a_models.append((idx, a_model))

    return a_models


def parse_atom_model(atom_model:AtomModel, disable_print=False) -> DataFrame:
    # 建立一个字典以备转化为DataFrame
    atom_info = {
        'Residue': [],      # str
        'Atom':    [],      # str
        'x':       [],      # float
        'y':       [],      # float
        'z':       [],      # float
        'R':       [],      # float
        'Type':    [],      # str
        'Surf':    [],      # float
        'Volu':    [],      # float
    }

    for atom in atom_model:
        Res = AA[atom[17:20]]
        Num = str(int(atom[22:26]))
        Ana = atom[12:16].strip()
        Ion = Ana[:2]
        
        # 残基、原子、坐标、半径、类型、表面积、体积
        atom_info['Residue'].append(atom[21] + Num + Res)
        atom_info['Atom']   .append(Ana)
        atom_info['x']      .append(float(atom[30:38]))
        atom_info['y']      .append(float(atom[38:46]))
        atom_info['z']      .append(float(atom[46:54]))

        if Res in Radius and Ana in Radius[Res]:
            __ = Radius[Res][Ana]
            atom_info['R']   .append(__[0])
            atom_info['Type'].append(__[1])
            atom_info['Surf'].append(__[2])
            atom_info['Volu'].append(__[3])
        
        # 单列Se和Cl是为了不与S和C重复
        elif Ion == 'Se' or Ion == 'Br' or Ion == 'Cl':
            __ = Radius['X'][Ion]
            atom_info['R']   .append(__[0])
            atom_info['Type'].append(__[1])
            atom_info['Surf'].append(__[2])
            atom_info['Volu'].append(__[3])
        
        # 排除离子，如果有的话直接删掉，因为很罕见所以不放在最前面判断，以免影响效率
        elif Ion in Metal:
            atom_info['Residue'] = atom_info['Residue'][:-1]
            atom_info['Atom']    = atom_info['Atom']   [:-1]
            atom_info['x']       = atom_info['x']      [:-1]
            atom_info['y']       = atom_info['y']      [:-1]
            atom_info['z']       = atom_info['z']      [:-1]
        
        # 如果不是双字母缩写的元素，那就是单字母的了
        else:
            __ = Radius['X'][Ana[0]]
            atom_info['R']   .append(__[0])
            atom_info['Type'].append(__[1])
            atom_info['Surf'].append(__[2])
            atom_info['Volu'].append(__[3])    
    
    atom_df = pd.DataFrame.from_dict(atom_info)
    if not disable_print:
        print(atom_df)
        print('len(atom_df):', len(atom_df))

    # NOTE: 加上水分子的半径
    atom_df['R'] += Radius_H2O

    for col in atom_df.columns:
        if is_numeric_dtype(atom_df[col]): # 把数值都变成np.float64类型
            atom_df[col] = atom_df[col].astype(DTYPE, copy=False)
    
    return atom_df


@timer_s
def arip_analyze(idx:int, name:str, atom_model:List[str], interval:float, ref_fp:Path, out_dp:Path, disable_print=False):
    # 先把二面角算出来
    pdb_file = StringIO(''.join(atom_model))
    dihedral_angle = pdb_dihedral_angle(name, pdb_file, disable_print)
    
    # 加载数据
    atom_df = parse_atom_model(atom_model, disable_print)

    # 不同的精度，不同的点数。计算接触体积
    contact_dict, volume = pdb_dotarray_volume(atom_df, interval, disable_print)
    
    # 删除不需要的列
    contact_df = atom_df[['Name', 'x', 'y', 'z', 'R', 'Surf', 'Type']]
    
    # 计算原子对之间的距离和接触面积
    surface = pdb_dotarray_surface(ref_fp, contact_df, contact_dict, disable_print)

    # 整理数据，并判断接触类型
    pdb_csv_sort(idx, name, dihedral_angle, surface, volume, out_dp, disable_print)
    

def arip_main(in_fp:Path, out_dp:Path, ref_fp:Path, interval:float, disable_print=False):
    name = Path(in_fp.stem).stem # 用两次stem是因为压缩文件还需要再去掉.pdb后缀
    a_models = load_atom_models(name, in_fp)

    try:
        for idx, a_model in a_models:
            t = arip_analyze(idx, name, a_model, interval, ref_fp, out_dp, disable_print)
            if idx == -1: print(f'The PDB {name} run OK, time cost: {t:.3f}s')
            else:         print(f'The PDB {name}_MODEL_{idx} run OK, time cost: {t:.3f}s')
    except:
        die(f'The PDB {name} cannot be analyzed, perhaps it contains IONS, or UNSUPPORTED format')
